import socket
import sys
import argparse
import tasks
import logging
import time

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
formatter = logging.Formatter('[%(levelname)s:%(asctime)s:%(module)s@%(lineno)s]:%(message)s',
                              datefmt='%Y-%m-%d %H:%M:%S')
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)

# register
def register(args,sock):
    task_list = tasks.load_tasks(args.config)
    for task in task_list:
        logger.info('acquiring GPU resource')
        timestamp = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
        taskid = hashlib.sha224(bytes(timestamp,'utf-8')).hexdigest()
        message = str({'action':'register','command':task.command,'memory':task.memory,'gpus':task.gpus,'id':taskid})
        sock.sendto(bytes(message,'utf-8'),(args.host,args.port))
        try:
            received = eval(str(sock.recv(4096),'utf-8').strip())
            response = received.get('response',None)
            if response == 'OK':
                gpu_id = received.get('ans')
                logger.info('run job {}@GPU-{}'.format(task.command,gpu_id))
                task.run(gpu_id)
                message = str({'action':'report','gpus':gpu_id})
                sock.sendto(bytes(message,'utf-8'),(args.host,args.port))
            elif response == 'ERROR':
                logger.error('ERROR:',received.get('ans','register failed'))
            elif response == 'CANCEL':
                logger.info('job {} cancelled'.format(task.command))
            else:
                logger.warning('unknown response')
        except KeyboardInterrupt:
            message = str({"action":"delete","multiple":False,'id':taskid})
            sock.sendto(bytes(message,'utf-8'),(args.host,args.port))
            received = str(sock.recv(4096),'utf-8').strip()

# query
def query(args,sock):
    message = str({"action":"query","idlen":args.id_len})
    sock.sendto(bytes(message,'utf-8'),(args.host,args.port))
    received = str(sock.recv(4096),'utf-8').strip()
    if received == '':
        received = 'no job in queue'
    else:
        received = 'jobs in queue:\n{}'.format(received)
    logger.info(received)

# delete
def delete(args,sock):
    assert args.id is not None
    message = str({"action":"delete","multiple":args.delete_multiple,'id':args.id})
    sock.sendto(bytes(message,'utf-8'),(args.host,args.port))
    received = str(sock.recv(4096),'utf-8').strip()

# move
def move(args,sock):
    assert args.id is not None
    assert args.position is not None
    message = str({"action":"move","id":args.id,"position":args.position})
    sock.sendto(bytes(message,'utf-8'),(args.host,args.port))


parser = argparse.ArgumentParser()
parser.add_argument('--port',type=int,default='10930')
parser.add_argument('--host',type=str,default='localhost')

subparsers = parser.add_subparsers()

# register
parser_register = subparsers.add_parser('register')
parser_register.set_defaults(callback=lambda args,sock:register(args,sock))
parser_register.add_argument('--config',type=str,required=True)

# query
parser_query = subparsers.add_parser('query')
parser_query.set_defaults(callback=lambda args,sock:query(args,sock))
parser_query.add_argument('--id-len',type=int,default=8)

# delete
parser_delete = subparsers.add_parser('delete')
parser_delete.set_defaults(callback=lambda args,sock:delete(args,sock))
parser_delete.add_argument('--id',type=str,required=True)
parser_delete.add_argument('--id-len',type=int,default=16)
parser_delete.add_argument('--delete-multiple',action='store_true',default=False)

# move
parser_move = subparsers.add_parser('move')
parser_move.set_defaults(callback=lambda args,sock:move(args,sock))
parser_move.add_argument('--id',type=str,required=True)
parser_move.add_argument('--position',type=int,required=True)


if __name__ == '__main__':
    sock = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)

    args = parser.parse_args()
    args.callback(args,sock)
